"""
重排模型评估工具模块 - 支持带分数的新格式
"""
from typing import List, Dict, Any, Union, Optional
from src.config import RELEVANCE_THRESHOLD

# 尝试导入必要的依赖库
try:
    import numpy as np
    import pandas as pd
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False
    print("警告: numpy 或 pandas 库未安装，评估功能将不可用。")
    print("请运行 'pip install numpy pandas' 安装所需依赖。")

def calculate_dcg(relevance_scores: List[float], k: Optional[int] = None) -> float:
    """
    计算DCG (Discounted Cumulative Gain)
    
    参数:
        relevance_scores: 相关性得分列表
        k: 取前k个元素计算，如果为None则使用全部
    
    返回:
        DCG值
    """
    if not NUMPY_AVAILABLE:
        print("无法计算DCG，缺少numpy库。")
        return 0.0
        
    if k is not None:
        relevance_scores = relevance_scores[:k]
    
    return np.sum([
        rel_score / np.log2(i + 2) 
        for i, rel_score in enumerate(relevance_scores)
    ])

def calculate_ndcg(relevance_scores: List[float], 
                  ideal_relevance_scores: List[float], 
                  k: Optional[int] = None) -> float:
    """
    计算NDCG (Normalized Discounted Cumulative Gain)
    
    参数:
        relevance_scores: 实际排序的相关性得分列表
        ideal_relevance_scores: 理想排序的相关性得分列表（已按降序排列）
        k: 取前k个元素计算，如果为None则使用全部
    
    返回:
        NDCG值，范围为[0,1]，越接近1表示排序质量越高
    """
    if not NUMPY_AVAILABLE:
        print("无法计算NDCG，缺少numpy库。")
        return 0.0
        
    dcg = calculate_dcg(relevance_scores, k)
    ideal_dcg = calculate_dcg(sorted(ideal_relevance_scores, reverse=True), k)
    
    if ideal_dcg == 0:
        return 0.0
    
    return dcg / ideal_dcg

def calculate_average_relevance(relevance_scores: List[float], k: Optional[int] = None) -> float:
    """
    计算平均相关性得分
    
    参数:
        relevance_scores: 相关性得分列表
        k: 取前k个元素计算，如果为None则使用全部
    
    返回:
        平均相关性得分
    """
    if not NUMPY_AVAILABLE:
        print("无法计算平均相关性，缺少numpy库。")
        return 0.0
        
    if k is not None:
        relevance_scores = relevance_scores[:k]
    
    if not relevance_scores:
        return 0.0
    
    return float(np.mean(relevance_scores))

def calculate_hit_rate(relevance_scores: List[float], threshold: float = RELEVANCE_THRESHOLD, k: Optional[int] = None) -> float:
    """
    计算命中率，即相关性得分大于等于阈值的比例
    
    参数:
        relevance_scores: 相关性得分列表
        threshold: 相关性阈值，大于等于该值被视为"有用"的知识片段
        k: 取前k个元素计算，如果为None则使用全部
    
    返回:
        命中率（0-1之间）
    """
    if k is not None:
        relevance_scores = relevance_scores[:k]
    
    if not relevance_scores:
        return 0.0
    
    hits = sum(1 for score in relevance_scores if score >= threshold)
    return hits / len(relevance_scores)

def evaluate_reranking_model(gold_standard_scores: Dict[str, List[float]], 
                            model_results: Dict[str, Dict],
                            k: int = 5,
                            verbose: bool = False) -> Dict[str, Dict[str, float]]:
    """
    评估重排模型性能
    
    参数:
        gold_standard_scores: 键为问题ID，值为每个知识片段的黄金标准相关性得分
        model_results: 键为问题ID，值为包含模型重排结果的字典
                      格式: {"question_id": {"scores": [score1, score2, ...], "rankings": [idx1, idx2, ...]}}
        k: 评估前k个结果
        verbose: 是否显示详细信息
    
    返回:
        评估结果字典，包含每个问题的NDCG@k、平均相关性得分、命中率和分数相关性
    """
    if not NUMPY_AVAILABLE:
        print("无法评估重排模型，缺少numpy库。")
        return {}
        
    results = {}
    total_questions = len(gold_standard_scores)
    processed = 0
    
    for question_id, gold_scores in gold_standard_scores.items():
        processed += 1
        
        if verbose and processed % 10 == 0:
            print(f"  进度: {processed}/{total_questions} ({processed/total_questions*100:.1f}%)")
            
        if question_id not in model_results:
            if verbose:
                print(f"  警告: 问题ID '{question_id}' 在模型结果中不存在，已跳过")
            continue
        
        # 获取排序和分数
        model_rankings = model_results[question_id]["rankings"]
        model_scores = model_results[question_id]["scores"]
        
        # 获取模型排序后的前k个知识片段索引
        model_top_k_indices = model_rankings[:k]
        
        # 获取这些知识片段的黄金标准相关性得分
        reranked_relevance = [gold_scores[idx] for idx in model_top_k_indices]
        
        # 计算评估指标
        ndcg = calculate_ndcg(reranked_relevance, gold_scores, k)
        avg_relevance = calculate_average_relevance(reranked_relevance)
        hit_rate = calculate_hit_rate(reranked_relevance, threshold=RELEVANCE_THRESHOLD)
        
        # 构建结果字典
        result_metrics = {
            f"ndcg@{k}": ndcg,
            f"avg_relevance@{k}": avg_relevance,
            f"hit_rate@{k}": hit_rate
        }
        
        # 如果verbose，打印详细信息
        if verbose and processed <= 5:  # 只打印前5个问题的详细信息
            print(f"\n  问题ID: {question_id}")
            print(f"    - NDCG@{k}: {ndcg:.4f}")
            print(f"    - 平均相关性@{k}: {avg_relevance:.4f}")
            print(f"    - 命中率@{k}: {hit_rate:.4f}")
            print(f"    - 模型Top-{k}排名: {model_top_k_indices}")
            print(f"    - 对应黄金标准得分: {reranked_relevance}")
        
        results[question_id] = result_metrics
    
    return results 

def test_ndcg_calculation():
    """
    测试NDCG计算的准确性，使用NDCG.md中的示例数据
    """
    print("\n" + "=" * 50)
    print("NDCG计算测试案例")
    print("=" * 50)
    
    # 案例1：基础示例
    # 模型排序后的相关性得分
    actual_relevance = [0.2, 1, 0.95, 1, 0.5]
    # 理想排序的相关性得分（降序排列）
    ideal_relevance = [1, 1, 0.95, 0.5, 0.2]
    
    # 手动计算DCG@5（使用实际的相关性得分）
    dcg_manual = (0.2/np.log2(2)) + (1.0/np.log2(3)) + (0.95/np.log2(4)) + (1.0/np.log2(5)) + (0.5/np.log2(6))
    # 手动计算IDCG@5（使用理想排序的相关性得分）
    idcg_manual = (1.0/np.log2(2)) + (1.0/np.log2(3)) + (0.95/np.log2(4)) + (0.5/np.log2(5)) + (0.2/np.log2(6))
    # 手动计算NDCG@5
    ndcg_manual = dcg_manual / idcg_manual
    
    # 使用函数计算DCG@5和NDCG@5
    dcg_func = calculate_dcg(actual_relevance, 5)
    ndcg_func = calculate_ndcg(actual_relevance, ideal_relevance, 5)
    
    print("案例1: 基础示例")
    print(f"实际排序的相关性得分: {actual_relevance}")
    print(f"理想排序的相关性得分: {ideal_relevance}")
    print(f"手动计算DCG@5: {dcg_manual:.4f}")
    print(f"函数计算DCG@5: {dcg_func:.4f}")
    print(f"手动计算IDCG@5: {idcg_manual:.4f}")
    print(f"手动计算NDCG@5: {ndcg_manual:.4f}")
    print(f"函数计算NDCG@5: {ndcg_func:.4f}")
    print(f"计算结果{'一致' if abs(ndcg_manual - ndcg_func) < 0.0001 else '不一致'}")
    
    # 案例2：更复杂的排序
    # 模型排序后的相关性得分
    actual_relevance2 = [0.3, 0.8, 0.1, 0.9, 0.5]
    # 理想排序的相关性得分（降序排列）
    ideal_relevance2 = [0.9, 0.8, 0.5, 0.3, 0.1]
    
    # 使用函数计算NDCG@5
    ndcg_func2 = calculate_ndcg(actual_relevance2, ideal_relevance2, 5)
    
    # 手动计算DCG@5（使用实际的相关性得分）
    dcg_manual2 = (0.3/np.log2(2)) + (0.8/np.log2(3)) + (0.1/np.log2(4)) + (0.9/np.log2(5)) + (0.5/np.log2(6))
    # 手动计算IDCG@5（使用理想排序的相关性得分）
    idcg_manual2 = (0.9/np.log2(2)) + (0.8/np.log2(3)) + (0.5/np.log2(4)) + (0.3/np.log2(5)) + (0.1/np.log2(6))
    # 手动计算NDCG@5
    ndcg_manual2 = dcg_manual2 / idcg_manual2
    
    print("\n案例2: 更复杂的排序")
    print(f"实际排序的相关性得分: {actual_relevance2}")
    print(f"理想排序的相关性得分: {ideal_relevance2}")
    print(f"手动计算DCG@5: {dcg_manual2:.4f}")
    print(f"手动计算IDCG@5: {idcg_manual2:.4f}")
    print(f"手动计算NDCG@5: {ndcg_manual2:.4f}")
    print(f"函数计算NDCG@5: {ndcg_func2:.4f}")
    print(f"计算结果{'一致' if abs(ndcg_manual2 - ndcg_func2) < 0.0001 else '不一致'}")
    
    # 案例3：测试使用rankings索引从gold_scores获取相关性得分的情况
    # 模拟gold_scores和model_rankings
    gold_scores = [3.0, 5.0, 2.0, 1.0, 4.0]
    model_rankings = [1, 4, 0, 2, 3]  # 排名顺序的索引，即排名第1的是索引1，排名第2的是索引4，...
    
    # 从gold_scores中获取按model_rankings排序后的相关性得分
    reranked_relevance = [gold_scores[idx] for idx in model_rankings]
    # 理想情况下的排序（gold_scores降序排列）
    ideal_relevance3 = sorted(gold_scores, reverse=True)
    
    # 使用函数计算NDCG@5
    ndcg_func3 = calculate_ndcg(reranked_relevance, gold_scores, 5)
    
    print("\n案例3: 使用rankings索引从gold_scores获取相关性得分")
    print(f"原始gold_scores: {gold_scores}")
    print(f"模型返回的rankings: {model_rankings}")
    print(f"重排后的相关性得分: {reranked_relevance}")
    print(f"理想排序的相关性得分: {ideal_relevance3}")
    print(f"函数计算NDCG@5: {ndcg_func3:.4f}")
    
    print("\n测试完成!")

# 如果直接运行此文件，则执行测试
if __name__ == "__main__":
    test_ndcg_calculation() 